import ast
import sys
import queue
import astpretty
from copy import deepcopy
from datetime import datetime
from colorama import init, Back, Style
from jinja2 import Environment, FileSystemLoader

init()
print(Style.RESET_ALL, end="")

CALL_SEQUENCE = []

EXECUTE_LINE_SEQOENCE = []

EXCEPTION_INFO_SEQUENCE = []

CONDITION_INFO_LIST = []

RECURRENT_INFO_LIST = []

CoverData = {} # filename -> set of line numbers

FILE_TO_TRACE = ""


def traverse_ast(ast_tree_root: ast.AST, focus_objs: list):
    interesting_objs = []
    # pre-order traverse
    temp_stack = []
    temp_stack.append(ast_tree_root)
    while len(temp_stack) > 0:
        temp_obj = temp_stack.pop()
        if type(temp_obj) in focus_objs:
            interesting_objs.append(temp_obj)
        if hasattr(temp_obj, 'body') is False:
            continue
        if hasattr(temp_obj, 'orelse') is True and temp_obj.orelse != []:
            temp_stack.extend(reversed(temp_obj.orelse))
        temp_stack.extend(reversed(temp_obj.body))

    # layer traverse
    # temp_q = queue.Queue()
    # temp_q.put(ast_tree_root)
    # while not temp_q.empty():
    #     temp_obj = temp_q.get()
    #     if type(temp_obj) in focus_objs:
    #         interesting_objs.append(temp_obj)
    #     if hasattr(temp_obj, 'body') is False:
    #         continue
    #     for obj in temp_obj.body:
    #         temp_q.put(obj)
    return interesting_objs


def trace(frame, event, arg):
    co = frame.f_code
    filename = co.co_filename
    lineno = frame.f_lineno
    if event == 'call':
        if filename == FILE_TO_TRACE:
            CALL_SEQUENCE.append((lineno, co.co_name))
    if event == 'exception' and filename == FILE_TO_TRACE:
        exc_type, exc_value, _ = arg
        EXCEPTION_INFO_SEQUENCE.append((exc_type.__name__, exc_value, lineno, co.co_name))
    if event == 'line':
        if filename == FILE_TO_TRACE:
            EXECUTE_LINE_SEQOENCE.append(lineno)
        line_dict = CoverData.get(filename, {})
        if lineno in line_dict:
            line_dict[lineno] += 1
        else:
            line_dict[lineno] = 1
        CoverData[filename] = line_dict
    return trace


def start():
    sys.settrace(trace)


def stop():
    sys.settrace(None)


def get_file_raw_lines(filename: str):
    with open(filename, 'r') as fread:
        raw_lines = [line for line in fread.readlines()]
    return raw_lines


def print_interesting_grammer_objs(interesting_objs: list):
    for obj in interesting_objs:
        print(type(obj))
        print(f"start line: {obj.lineno}, end line: {obj.end_lineno}")
        if isinstance(obj, ast.If):
            if obj.orelse != []:
                print(obj.orelse[0].lineno)
        print('---------')


def compress_executed_sequence(raw_executed_sequence: list):
    executed_sequence = deepcopy(raw_executed_sequence)
    executed_sequence.append(-1)  # fake tail
    compressed_executed_sequence = []
    sequence_length = len(executed_sequence)
    counter = 0
    while counter < sequence_length:
        if counter == 0:
            prev_lineno = executed_sequence[counter]
            prev_lineno_executed_amount = 1
        else:
            if executed_sequence[counter] == prev_lineno:
                prev_lineno_executed_amount += 1
            else:
                compressed_executed_sequence.append((str(prev_lineno).zfill(3), prev_lineno_executed_amount))
                prev_lineno = executed_sequence[counter]
                prev_lineno_executed_amount = 1
        counter += 1
    # tuple(lineno, amount_of_executed)
    return compressed_executed_sequence


def report(filename: str, is_colorful: bool=False) -> tuple:
    html_show_code_list = []
    color_prefix = '' if is_colorful is False else Back.GREEN 
    blank_line = 0
    covered = CoverData[filename]
    ncovered = len(covered)
    nline = 0

    print('\nFilename: %s' % filename)
    with open(filename, 'r', encoding='utf8') as f:
        for line in f.readlines():
            line = line.rstrip()
            nline += 1
            if line.strip() == '' or line.strip().startswith('#') or line.strip().startswith('\'\'\'') or line.strip().startswith('\"\"\"'):
                line_content = 'Line %s| %s' % (str(nline).zfill(3), line)
                html_show_code_list.append(('rgb(5,67,238)', ' '+line_content, 0))
                print(Back.CYAN+line_content, end='')
                blank_line += 1
            else:
                if nline in covered:
                    color_prefix = Back.GREEN
                    color_str = 'rgb(14,223,14)'
                    exec_times = covered[nline]
                else:
                    color_prefix = Back.RED
                    color_str = 'rgb(240,56,56)'
                    exec_times = 0
                line_content = "Line %s| %s" % (str(nline).zfill(3), line)
                html_show_code_list.append((color_str, ' '+line_content, exec_times))
                print(color_prefix+"Line %s| %s" % (str(nline).zfill(3), line), end='')
                # print('--------------------------')
            print(Style.RESET_ALL, end='')
            print()
    print(Style.RESET_ALL, end='')

    nline -= blank_line
    print(f"\nCode lines: {nline}, covered lines: {ncovered}, blank or comment lines: {blank_line}")
    print("Code coverage: %.2f%%" % (100 * float(ncovered)/nline))
    print(f"Line executed sequence: ", EXECUTE_LINE_SEQOENCE)
    print(f"Function Call line: ", CALL_SEQUENCE[1:])
    if len(EXCEPTION_INFO_SEQUENCE) == 0:
        print("None exception")
    else:
        print('* Tracing exception:\n'
             f'* {EXCEPTION_INFO_SEQUENCE[0][0]} "{EXCEPTION_INFO_SEQUENCE[0][1]}"\n'
             f'* on line {EXCEPTION_INFO_SEQUENCE[0][2]} of {EXCEPTION_INFO_SEQUENCE[0][3]}\n'
        )
    return html_show_code_list, nline, blank_line, ncovered, 100 * float(ncovered)/nline


def generate_html_report(**html_info_dict):
    env = Environment(loader=FileSystemLoader('template'))

    temp = env.get_template('index.html')
    html_content = temp.render(**html_info_dict)
    with open('test_res.html', 'w', encoding='utf8') as fw:
        fw.write(html_content)


def analyse_interseting_grammer_objs(interesting_objs: list):
    global CONDITION_INFO_LIST, RECURRENT_INFO_LIST
    for obj in interesting_objs:
        if isinstance(obj, ast.For) or isinstance(obj, ast.While):
            RECURRENT_INFO_LIST.append((obj.lineno, obj.end_lineno))
        elif isinstance(obj, ast.If):
            else_lineno = -1
            if obj.orelse != []:
                else_lineno = obj.orelse[0].lineno
            CONDITION_INFO_LIST.append((obj.lineno, obj.end_lineno, else_lineno))


if __name__ == '__main__':
    if len(sys.argv) == 1:
        sys.exit(0)
    
    filename, args = sys.argv[1], sys.argv[2:]
    FILE_TO_TRACE = filename
    try:
        with open(filename, 'r', encoding='utf8') as f:
            content = f.read()
        ast_node = ast.parse(content, "msg.log", mode="exec")
        interesting_objs = traverse_ast(ast_node, focus_objs=[ast.FunctionDef, ast.For, ast.While, ast.If])
        analyse_interseting_grammer_objs(interesting_objs)
        # print(interesting_objs)
        # astpretty.pprint(ast_node)
        code = compile(content, filename, 'exec')
    except IOError as ioe:
        print('No such file: %s' % filename)
        sys.exit(1)

    start()
    try:
        print(f"-----File {FILE_TO_TRACE}'s output"+"-----")
        exec(code)
        print("------------------------------------")
    finally:
        stop()
        print(f"\n-----File {FILE_TO_TRACE}'s report"+"-----")
        html_show_code_list, code_lines, blank_and_comment_lines, executed_lines, code_coverage_rate = report(filename, True)
        generate_html_report(**{
            'filename': filename,
            'tested_code': html_show_code_list,
            'code_lines': code_lines,
            'blank_and_comment_lines': blank_and_comment_lines,
            'executed_lines': executed_lines,
            'code_coverage_rate': '%.2f'%code_coverage_rate,
            'recurrent_info_list': [(html_show_code_list[start_lineno-1: end_lineno], start_lineno, end_lineno) for start_lineno, end_lineno in RECURRENT_INFO_LIST],
            'condition_info_list': [(html_show_code_list[start_lineno-1: end_lineno], start_lineno, end_lineno, else_lineno) for start_lineno, end_lineno, else_lineno in  CONDITION_INFO_LIST],
            'compressed_executed_sequence': compress_executed_sequence(EXECUTE_LINE_SEQOENCE),
            'call_sequence': [(str(lineno).zfill(3), each_call.replace('<', '(').replace('>', ')')) for lineno, each_call in CALL_SEQUENCE[1:]],
            'exception_info_sequence': [(i, j, k, l) for i, j, k, l in EXCEPTION_INFO_SEQUENCE],
            'datetime_str': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        })
        print("------------------------------------")
        sys.exit(0)
